/**
 * ComfyUI Llama-CPP All-In-One Frontend Extension
 * Adds Generate/Stop buttons and result viewer
 */

import { app } from "../../scripts/app.js";

// Extension name
const EXTENSION_NAME = "BuimenLabo_LlamaCppAIO";

// Create button widget
function createControlButton(node, label, className, callback) {
    const button = document.createElement("button");
    button.textContent = label;
    button.className = `comfy-button ${className}`;
    button.style.cssText = `
        width: 100px;
        height: 35px;
        border: none;
        border-radius: 4px;
        cursor: pointer;
        font-size: 12px;
        font-weight: bold;
        margin: 2px;
        transition: all 0.2s;
    `;
    
    // Set button colors based on type
    if (className.includes('generate')) {
        button.style.background = "#5cb85c";
        button.style.color = "white";
    } else if (className.includes('reset')) {
        button.style.background = "#f0ad4e";
        button.style.color = "white";
    } else if (className.includes('unload')) {
        button.style.background = "#d9534f";
        button.style.color = "white";
    }
    
    button.addEventListener("click", callback);
    button.addEventListener("mouseenter", () => {
        if (!button.disabled) {
            if (className.includes('generate')) {
                button.style.background = "#6cc76c";
            } else if (className.includes('reset')) {
                button.style.background = "#f6c969";
            } else if (className.includes('unload')) {
                button.style.background = "#e9645f";
            }
        }
    });
    button.addEventListener("mouseleave", () => {
        if (!button.disabled) {
            if (className.includes('generate')) {
                button.style.background = "#5cb85c";
            } else if (className.includes('reset')) {
                button.style.background = "#f0ad4e";
            } else if (className.includes('unload')) {
                button.style.background = "#d9534f";
            }
        }
    });
    
    return button;
}

// Create result viewer widget
function createResultViewer(node) {
    const container = document.createElement("div");
    container.className = "result-viewer-container";
    container.style.cssText = `
        width: 100%;
        max-width: 400px;
        margin: 5px 0;
        border: 1px solid #555;
        border-radius: 4px;
        background: #1a1a1a;
    `;
    
    // Header
    const header = document.createElement("div");
    header.className = "result-viewer-header";
    header.style.cssText = `
        padding: 8px 12px;
        background: #333;
        border-bottom: 1px solid #555;
        display: flex;
        align-items: center;
        justify-content: space-between;
    `;
    
    const title = document.createElement("span");
    title.textContent = "📺 生成結果ビューワ";
    title.style.cssText = `
        color: white;
        font-size: 12px;
        font-weight: bold;
    `;
    
    const status = document.createElement("span");
    status.textContent = "待機中";
    status.className = "viewer-status";
    status.style.cssText = `
        color: #ccc;
        font-size: 11px;
        font-family: monospace;
    `;
    
    header.appendChild(title);
    header.appendChild(status);
    
    // Content area
    const content = document.createElement("div");
    content.className = "result-viewer-content";
    content.style.cssText = `
        padding: 10px;
        min-height: 80px;
        max-height: 200px;
        overflow-y: auto;
        color: white;
        font-family: 'Courier New', monospace;
        font-size: 12px;
        line-height: 1.4;
        white-space: pre-wrap;
    `;
    content.textContent = "生成結果がここに表示されます...";
    
    container.appendChild(header);
    container.appendChild(content);
    
    return { container, content, status };
}

// Enhanced node creation for AIO
function enhanceAIONode(node) {
    console.log("🔧 Enhancing Llama-CPP AIO node:", node.comfyClass);
    
    // Add button container (5行目: [Generate][Stop] ボタン)
    const buttonContainer = document.createElement("div");
    buttonContainer.className = "aio-button-container";
    buttonContainer.style.cssText = `
        display: flex;
        gap: 10px;
        margin: 5px 0;
        justify-content: center;
    `;
    
    // Generate button
    const generateButton = createControlButton(
        node,
        "🚀 Generate",
        "aio-generate-button",
        () => executeGenerate(node)
    );
    
    // Reset button
    const resetButton = createControlButton(
        node,
        "🔄 Reset",
        "aio-reset-button",
        () => executeReset(node)
    );
    
    // Unload button
    const unloadButton = createControlButton(
        node,
        "🗑️ Unload",
        "aio-unload-button",
        () => executeUnload(node)
    );
    
    buttonContainer.appendChild(generateButton);
    buttonContainer.appendChild(resetButton);
    buttonContainer.appendChild(unloadButton);
    
    // Add buttons to node
    node.addDOMWidget("aio_buttons", "div", buttonContainer);
    
    // Add result viewer (4行目: 📺 生成結果ビューワ)
    const viewer = createResultViewer(node);
    node.addDOMWidget("result_viewer", "div", viewer.container);
    
    // Store references for easy access
    node.generateButton = generateButton;
    node.resetButton = resetButton;
    node.unloadButton = unloadButton;
    node.resultViewer = viewer.content;
    node.resultStatus = viewer.status;
    
    // Watch for widget changes
    const originalOnWidgetChange = node.onWidgetChange;
    node.onWidgetChange = function(widget) {
        if (originalOnWidgetChange) {
            originalOnWidgetChange.call(this, widget);
        }
        if (widget.name === "generated_result") {
            updateResultViewer(this, widget.value);
            console.log("📺 Generated result updated:", widget.value?.slice(0, 100));
        }
    };
    
    console.log("✅ AIO buttons and viewer added to node");
}

// Execute generate action (Direct API call like prompt translator)
async function executeGenerate(node) {
    console.log("🚀 Generate button clicked!");
    
    // Update UI state
    disableButtons(node, true);
    updateStatus(node, "生成中...", "#4a90e2");
    updateResultViewer(node, "🚀 生成開始: パラメータ収集中...");
    
    // Get all widget values
    const instructionWidget = node.widgets.find(w => w.name === "instruction_prompt");
    const inputWidget = node.widgets.find(w => w.name === "input_text");
    const modelWidget = node.widgets.find(w => w.name === "model_path");
    const maxTokensWidget = node.widgets.find(w => w.name === "max_tokens");
    const temperatureWidget = node.widgets.find(w => w.name === "temperature");
    const topKWidget = node.widgets.find(w => w.name === "top_k");
    const topPWidget = node.widgets.find(w => w.name === "top_p");
    const repeatPenaltyWidget = node.widgets.find(w => w.name === "repeat_penalty");
    const nCtxWidget = node.widgets.find(w => w.name === "n_ctx");
    const nGpuLayersWidget = node.widgets.find(w => w.name === "n_gpu_layers");
    const nsfwFilterWidget = node.widgets.find(w => w.name === "nsfw_filter");
    const autoUnloadWidget = node.widgets.find(w => w.name === "auto_unload");
    
    const instructionPrompt = instructionWidget ? instructionWidget.value : "";
    const inputText = inputWidget ? inputWidget.value : "";
    const modelPath = modelWidget ? modelWidget.value : "";
    const maxTokens = maxTokensWidget ? maxTokensWidget.value : 128;
    const temperature = temperatureWidget ? temperatureWidget.value : 0.8;
    const topK = topKWidget ? topKWidget.value : 40;
    const topP = topPWidget ? topPWidget.value : 0.95;
    const repeatPenalty = repeatPenaltyWidget ? repeatPenaltyWidget.value : 1.1;
    const nCtx = nCtxWidget ? nCtxWidget.value : 2048;
    const nGpuLayers = nGpuLayersWidget ? nGpuLayersWidget.value : 40;
    const nsfwFilter = nsfwFilterWidget ? nsfwFilterWidget.value : true;
    const autoUnload = autoUnloadWidget ? autoUnloadWidget.value : true;
    
    if (!inputText.trim() && !instructionPrompt.trim()) {
        alert("入力テキストまたは指示プロンプトを入力してください");
        disableButtons(node, false);
        updateStatus(node, "エラー: 入力なし", "#d9534f");
        return;
    }
    
    console.log(`🔄 Starting generation with model: ${modelPath}`);
    updateResultViewer(node, "📡 APIに接続中...");
    
    try {
        // Call direct API (like prompt translator)
        const result = await callLlamaAPI(
            instructionPrompt, inputText, modelPath, maxTokens, temperature, 
            topK, topP, repeatPenalty, nCtx, nGpuLayers, nsfwFilter, autoUnload
        );
        
        console.log(`✅ Generation completed: "${result.slice(0, 100)}..."`);
        
        // Update result viewer and widget
        updateResultViewer(node, result);
        updateStatus(node, "生成完了", "#5cb85c");
        
        // Update the generated_result widget
        const resultWidget = node.widgets.find(w => w.name === "generated_result");
        if (resultWidget) {
            resultWidget.value = result;
            console.log("📝 Updated generated_result widget");
        }
        
    } catch (error) {
        console.error(`❌ Generation failed:`, error);
        updateResultViewer(node, `エラー: ${error.message}`);
        updateStatus(node, "エラー", "#d9534f");
    } finally {
        disableButtons(node, false);
    }
}

// Call Llama API directly (like prompt translator)
async function callLlamaAPI(instructionPrompt, inputText, modelPath, maxTokens, temperature, 
                           topK, topP, repeatPenalty, nCtx, nGpuLayers, nsfwFilter, autoUnload) {
    try {
        console.log(`🌐 API call: ${modelPath} "${inputText.slice(0, 50)}..."`);
        
        // Use FormData as required by ComfyUI
        const formData = new FormData();
        formData.append('instruction_prompt', instructionPrompt);
        formData.append('input_text', inputText);
        formData.append('model_path', modelPath);
        formData.append('max_tokens', maxTokens.toString());
        formData.append('temperature', temperature.toString());
        formData.append('top_k', topK.toString());
        formData.append('top_p', topP.toString());
        formData.append('repeat_penalty', repeatPenalty.toString());
        formData.append('n_ctx', nCtx.toString());
        formData.append('n_gpu_layers', nGpuLayers.toString());
        formData.append('nsfw_filter', nsfwFilter.toString());
        formData.append('auto_unload', autoUnload.toString());
        
        const response = await fetch('/llama_cpp/generate', {
            method: 'POST',
            body: formData
        });
        
        console.log(`📡 API response status: ${response.status}`);
        
        if (!response.ok) {
            const errorText = await response.text();
            console.error(`❌ API error response:`, errorText);
            throw new Error(`HTTP ${response.status}: ${errorText}`);
        }
        
        const result = await response.json();
        console.log(`📋 API response:`, result);
        
        if (!result.success) {
            throw new Error(result.error || 'Generation failed');
        }
        
        return result.generated_text;
        
    } catch (error) {
        console.error(`💥 API error:`, error);
        if (error.name === 'TypeError' && error.message.includes('fetch')) {
            throw new Error(`API not available. Please restart ComfyUI.`);
        }
        throw error;
    }
}

// Execute reset action
function executeReset(node) {
    console.log("🔄 Reset button clicked!");
    
    // Reset all widgets to default values
    const widgets = {
        instruction_prompt: "You are a prompt generator for Stable Diffusion.\n\nTask:\nConvert the input description into ONE optimized English prompt for Stable Diffusion.\nRules:\n- Output format: single line of comma-separated phrases only\n- Stay faithful to the original description's content and intent\n- Include: subject, setting/environment, lighting, mood/atmosphere, artistic style, quality tags\n- Use descriptive English terms that accurately represent the input\n- Do not censor or sanitize the content - convert faithfully\n- End cleanly after final quality tag\n\nOutput:\n<prompt>",
        input_text: "プールサイドで寝そべっている女性",
        generated_result: "",
        max_tokens: 128,
        temperature: 0.8,
        top_k: 40,
        top_p: 0.95,
        repeat_penalty: 1.1,
        use_random_seed: true,
        seed: 0,
        nsfw_filter: true,
        auto_unload: true
    };
    
    // Apply default values to widgets
    for (const [widgetName, defaultValue] of Object.entries(widgets)) {
        const widget = node.widgets.find(w => w.name === widgetName);
        if (widget) {
            widget.value = defaultValue;
            console.log(`📝 Reset ${widgetName} to default: ${defaultValue}`);
        }
    }
    
    // Reset UI state
    updateStatus(node, "初期値に戻しました", "#5cb85c");
    updateResultViewer(node, "生成結果がここに表示されます...");
    
    // Re-enable buttons
    disableButtons(node, false);
    
    console.log("🔄 Reset to defaults completed");
}

// Execute unload action
async function executeUnload(node) {
    console.log("🗑️ Unload button clicked!");
    
    // Update UI state
    disableButtons(node, true);
    updateStatus(node, "VRAM解放中...", "#d9534f");
    updateResultViewer(node, "🗑️ モデルをVRAMから解放しています...");
    
    try {
        // Call unload API
        const response = await fetch('/llama_cpp/unload', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({})
        });
        
        if (response.ok) {
            const result = await response.json();
            console.log("✅ Unload successful:", result);
            
            updateStatus(node, "VRAM解放完了", "#5cb85c");
            updateResultViewer(node, "✅ VRAMから正常に解放されました。\n\n" + (result.message || "メモリが解放されました。"));
            
            // Show memory status after 2 seconds
            setTimeout(() => {
                updateResultViewer(node, "✅ VRAM解放完了\n\n💡 VRAM使用量の確認:\n• タスクマネージャー (GPU使用量)\n• nvidia-smi コマンド\n• ComfyUIコンソールログを確認");
            }, 2000);
            
        } else {
            throw new Error(`HTTP ${response.status}: ${await response.text()}`);
        }
        
    } catch (error) {
        console.error("❌ Unload failed:", error);
        updateStatus(node, "解放エラー", "#d9534f");
        updateResultViewer(node, `❌ VRAM解放エラー: ${error.message}`);
    } finally {
        disableButtons(node, false);
    }
}

// Update result viewer content
function updateResultViewer(node, content) {
    if (node.resultViewer) {
        node.resultViewer.textContent = content || "生成結果がここに表示されます...";
        console.log("📺 Updated result viewer");
    }
}

// Update status display
function updateStatus(node, status, color) {
    if (node.resultStatus) {
        node.resultStatus.textContent = status;
        node.resultStatus.style.color = color || "#ccc";
    }
}

// Disable/enable buttons
function disableButtons(node, disabled) {
    if (node.generateButton) {
        node.generateButton.disabled = disabled;
        node.generateButton.style.opacity = disabled ? "0.6" : "1";
        node.generateButton.style.cursor = disabled ? "not-allowed" : "pointer";
    }
    
    if (node.resetButton) {
        node.resetButton.disabled = disabled;
        node.resetButton.style.opacity = disabled ? "0.6" : "1";
        node.resetButton.style.cursor = disabled ? "not-allowed" : "pointer";
    }
    
    if (node.unloadButton) {
        node.unloadButton.disabled = disabled;
        node.unloadButton.style.opacity = disabled ? "0.6" : "1";
        node.unloadButton.style.cursor = disabled ? "not-allowed" : "pointer";
    }
}

// Register extension
app.registerExtension({
    name: EXTENSION_NAME,
    
    async nodeCreated(node) {
        console.log("🔍 Node created:", node.comfyClass);
        
        // Support both original and BuimenLabo versions
        if (node.comfyClass === "LlamaCppAIO" || node.comfyClass === "BuimenLabo_LlamaCppAIO") {
            console.log("🚀 Enhancing Llama-CPP AIO node:", node.comfyClass);
            enhanceAIONode(node);
        }
    },
    
    async setup() {
        console.log("🚀 Llama-CPP AIO extension setup");
        
        // Add custom CSS
        const style = document.createElement("style");
        style.textContent = `
            .aio-button-container {
                display: flex;
                gap: 10px;
                margin: 5px 0;
                justify-content: center;
            }
            
            .aio-generate-button:hover:not(:disabled) {
                transform: translateY(-1px);
                box-shadow: 0 2px 4px rgba(0,0,0,0.2);
            }
            
            .aio-reset-button:hover:not(:disabled) {
                transform: translateY(-1px);
                box-shadow: 0 2px 4px rgba(0,0,0,0.2);
            }
            
            .aio-unload-button:hover:not(:disabled) {
                transform: translateY(-1px);
                box-shadow: 0 2px 4px rgba(0,0,0,0.2);
            }
            
            .aio-generate-button:active {
                transform: translateY(1px);
            }
            
            .aio-reset-button:active {
                transform: translateY(1px);
            }
            
            .aio-unload-button:active {
                transform: translateY(1px);
            }
            
            .result-viewer-container {
                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            }
            
            .result-viewer-content::-webkit-scrollbar {
                width: 6px;
            }
            
            .result-viewer-content::-webkit-scrollbar-track {
                background: #2a2a2a;
            }
            
            .result-viewer-content::-webkit-scrollbar-thumb {
                background: #555;
                border-radius: 3px;
            }
            
            .result-viewer-content::-webkit-scrollbar-thumb:hover {
                background: #777;
            }
        `;
        
        document.head.appendChild(style);
        console.log("📦 AIO CSS loaded");
    }
});

console.log("📋 BuimenLabo Llama-CPP AIO extension loaded successfully!");
console.log("🔧 Extension file: builmenlabo_llama.js");